# encoding=utf-8
import asyncio
import base64
import hashlib
import json
import sys
import time
import aioredis
from loguru import logger
from allure_commons.types import Severity
from playwright.async_api import async_playwright, Page, Playwright
from config import WebSocketToken

logger.remove()  # 移除默认的日志处理器
logger.add("log/app.log", level="INFO", rotation="10 MB", retention="30 days", compression="zip")  # 输出到文件
logger.add(sys.stderr, level="INFO")  # 输出到控制台


class WebAutomation:

    def __init__(self, headless=True, proxy=None, viewport=None, user_agent=None):
        """
        初始化 WebAutomation 实例

        :param headless: 默认为 True (无头模式)，如果设置为 False，则浏览器会以窗口形式打开
        :param proxy: 代理地址，格式为 'http://ip:port' 或 {'server': 'http://ip:port', 'username':'用户名', 'password':'密码'}
        """
        self.headless = headless
        self.proxy = proxy
        self.viewport = viewport
        self.user_agent = user_agent
        self.context = None
        self.page = None
        self.browser = None

    async def start_browser(self, playwright: Playwright):
        """启动浏览器"""
        browser_args = {
            "headless": self.headless,
            "args": [
                '--no-sandbox',
                '--disable-setuid-sandbox',
                '--disable-dev-shm-usage',
                '--disable-accelerated-2d-canvas',
                '--disable-gpu',
                '--disable-extensions',
                '--disable-web-security',
                '--disable-features=site-per-process',
                '--disable-features=IsolateOrigins',
                '--disable-site-isolation-trials'
            ]
        }
        if self.proxy:
            if isinstance(self.proxy, str):
                browser_args["proxy"] = {"server": self.proxy}
            elif isinstance(self.proxy, dict):
                browser_args["proxy"] = self.proxy

        self.browser = await playwright.chromium.launch(**browser_args)
        self.context = await self.browser.new_context(viewport=self.viewport, user_agent=self.user_agent)
        self.page = await self.context.new_page()
        self.page.set_default_timeout(30000)

    async def close_browser(self):
        """关闭浏览器"""
        if self.context:
            await self.context.close()
        if self.browser:
            await self.browser.close()

    # 跳转到指定页面
    async def navigate_to(self, url, page_width=None, page_height=None):
        if page_width and page_height:
            await self.page.set_viewport_size({"width": page_width, "height": page_height})
        await self.page.goto(url, wait_until="domcontentloaded")  # 等待 DOM 内容加载完成
        await self.page.wait_for_load_state("networkidle")  # 等待网络请求完成，确保页面稳定
        return self.page

    # 打开第二个页面
    async def open_new_page(self, url, width=None, height=None):
        # 在同一个浏览器上下文中创建新页面
        new_page = await self.page.context.new_page()
        if width and height:
            await new_page.set_viewport_size({"width": width, "height": height})
        await new_page.goto(url, wait_until="domcontentloaded")
        await self.page.wait_for_load_state("networkidle")
        return new_page

    # 关闭当前页面
    async def close_current_page(self):
        await self.page.close()

    # 点击
    async def click(self, selector):
        await self.page.locator(selector).click()

    # 悬停
    async def hover(self, selector):
        await self.page.locator(selector).hover()

    # 输入
    async def input(self, selector, text):
        await self.page.locator(selector).click()
        await self.page.locator(selector).fill(text)

    # 刷新
    async def refresh(self):
        await self.page.reload()
        await self.page.wait_for_load_state("networkidle")

    # 等待
    async def wait(self, time_ms: int):
        await self.page.wait_for_timeout(time_ms)

    # 截图
    async def screenshot(self):
        """截图方法"""
        try:
            await self.page.wait_for_load_state('domcontentloaded', timeout=3000)
            return await self.page.screenshot(timeout=5000)
        except Exception as e:
            print(f"截图时发生错误: {e}")
            await self.page.wait_for_load_state('domcontentloaded', timeout=3000)
            return await self.page.screenshot(timeout=5000)

    # alert
    async def alert(self, accept=None):
        """
        :param accept: 有值点击确定，无值点击取消
        :param page: 指定页面，默认为当前页面
        """
        # 监听dialog事件
        dialog = await self.page.wait_for_event("dialog")
        if dialog:
            if accept:
                await dialog.accept()
            else:
                await dialog.dismiss()

    # 获取元素文本
    async def get_text(self, selector):
        return await self.page.locator(selector).inner_text()

    # 下拉选项框
    async def select(self, selector, option_text):
        """
        :param selector: 例：页面结构#dropdown：<select id="dropdown" name="options">
        :param option_text: 选项文本
        :param page: 指定页面，默认为当前页面
        """
        await self.page.locator(f"select{selector}").select_option(option_text)

    # iframe
    async def switch_to_iframe(self, frame_locator: list, locator):
        """
        :param frame_locator: 传列表，iframe定位，支持id name class css xpath text等定位方式
        :param locator: 元素定位
        :param page: 指定页面，默认为当前页面
        """
        # 初始页面的frame
        frame = self.page

        # 逐层定位 iframe
        for locator_item in frame_locator:
            # 使用 await 来确保每一层 iframe 都能正确获取
            frame = await frame.frame_locator(locator_item)

        # 获取目标元素并点击
        await frame.locator(locator).click()

    async def parent_element_positioning_iframe_click(self, parent_element, layer: int, child_iframe_index: list,
                                                      locator):
        """
        :param layer: 层嵌套，1为第一层，2为第二层
        :param parent_element: 父级元素定位，支持id name class css xpath text等定位方式
        :param child_iframe_index: 子iframe索引，从0开始计数
        :param locator: 元素定位
        :param page: 指定页面，默认为当前页面
        """
        page = self.page
        # 直接进入iframe
        if layer == 0:
            iframe = await page.query_selector(parent_element)
            frame = await iframe.content_frame()
            await frame.locator(locator).click()
        # 通过父元素进入第一层iframe
        elif layer == 1:
            iframe = await page.query_selector(parent_element)
            frame = await (await iframe.query_selector('iframe')).content_frame()
            await frame.locator(locator).click()
        # 通过父元素进入第二层iframe
        elif layer == 2:
            iframe = await page.query_selector(parent_element)
            frame = await (await iframe.query_selector('iframe')).content_frame()
            child_iframe = await (await frame.query_selector_all('iframe'))[int(child_iframe_index[0])].content_frame()
            await child_iframe.locator(locator).click()
        # 通过父元素进入第一层iframe
        elif layer == 3:
            iframe = await page.query_selector(parent_element)
            frame = await (await iframe.query_selector('iframe')).content_frame()
            child_iframe = await (await frame.query_selector_all('iframe'))[int(child_iframe_index[0])].content_frame()
            grandchild_iframe = await (await child_iframe.query_selector_all('iframe'))[
                int(child_iframe_index[1])].content_frame()
            await grandchild_iframe.locator(locator).click()
        else:
            raise Exception("layer参数错误，目前仅支持三层嵌套")

    async def parent_element_positioning_iframe_input(self, layer: int, parent_element, child_iframe_index: list,
                                                      locator,
                                                      text):
        """
        :param layer: 层嵌套，1为第一层，2为第二层
        :param text: 输入文本
        :param parent_element: 父级元素定位，支持id name class css xpath text等定位方式
        :param child_iframe_index: 子iframe索引，从0开始计数
        :param locator: 元素定位
        :param page: 指定页面，默认为当前页面
        """
        page = self.page
        # 直接进入iframe
        if layer == 0:
            iframe = await page.query_selector(parent_element)
            frame = await iframe.content_frame()
            await frame.locator(locator).click()
            await frame.locator(locator).fill(text)
        # 通过父元素进入第一层iframe
        elif layer == 1:
            iframe = await page.query_selector(parent_element)
            frame = await (await iframe.query_selector('iframe')).content_frame()
            await frame.locator(locator).click()
            await frame.locator(locator).fill(text)
        # 通过父元素进入第二层iframe
        elif layer == 2:
            iframe = await page.query_selector(parent_element)
            frame = await (await iframe.query_selector('iframe')).content_frame()
            child_iframe = await (await frame.query_selector_all('iframe'))[int(child_iframe_index[0])].content_frame()
            await child_iframe.locator(locator).click()
            await child_iframe.locator(locator).fill(text)
        # 通过父元素进入第三层iframe
        elif layer == 3:
            iframe = await page.query_selector(parent_element)
            frame = await (await iframe.query_selector('iframe')).content_frame()
            child_iframe = await (await frame.query_selector_all('iframe'))[int(child_iframe_index[0])].content_frame()
            grandchild_iframe = await (await child_iframe.query_selector_all('iframe'))[
                int(child_iframe_index[1])].content_frame()
            await grandchild_iframe.locator(locator).click()
            await grandchild_iframe.locator(locator).fill(text)
        else:
            raise Exception("layer参数错误，目前仅支持三层嵌套")

    async def parent_element_positioning_iframe_select(self, parent_element, layer: int, child_iframe_index: list,
                                                       selector, option_text):
        """
        :param layer: 层嵌套，1为第一层，2为第二层
        :param parent_element: 父级元素定位，支持id name class css xpath text等定位方式
        :param child_iframe_index: 子iframe索引，从0开始计数
        :param page: 指定页面，默认为当前页面
        :param selector: 例：页面结构#dropdown：<select id="dropdown" name="options">
        :param option_text: 选项文本
        """
        page = self.page
        # 直接进入iframe
        if layer == 0:
            iframe = await page.query_selector(parent_element)
            frame = await iframe.content_frame()
            await frame.locator(f"select{selector}").select_option(option_text)
        # 通过父元素进入第一层iframe
        elif layer == 1:
            iframe = await page.query_selector(parent_element)
            frame = await (await iframe.query_selector('iframe')).content_frame()
            await frame.locator(f"select{selector}").select_option(option_text)
        # 通过父元素进入第二层iframe
        elif layer == 2:
            iframe = await page.query_selector(parent_element)
            frame = await (await iframe.query_selector('iframe')).content_frame()
            child_iframe = await (await frame.query_selector_all('iframe'))[int(child_iframe_index[0])].content_frame()
            await child_iframe.locator(f"select{selector}").select_option(option_text)
        # 第通过父元素进入第三层iframe
        elif layer == 3:
            iframe = await page.query_selector(parent_element)
            frame = await (await iframe.query_selector('iframe')).content_frame()
            child_iframe = await (await frame.query_selector_all('iframe'))[int(child_iframe_index[0])].content_frame()
            grandchild_iframe = await (await child_iframe.query_selector_all('iframe'))[
                int(child_iframe_index[1])].content_frame()
            await grandchild_iframe.locator(f"select{selector}").select_option(option_text)
        else:
            raise Exception("layer参数错误，目前仅支持三层嵌套")

    # 获取当前窗口URL
    async def get_url(self):
        return self.page.url

    # 获取当前窗口标题
    async def get_title(self):
        return await self.page.title()

    # 获取Cookie
    async def get_cookie(self):
        return await self.page.context.cookies()

    # 写入Cookie
    async def set_cookie(self, cookie: list):
        await self.page.context.add_cookies(cookie)

    # 删除Cookie
    async def delete_cookie(self):
        await self.page.context.clear_cookies()

    # 获取当前页面源码
    async def get_html(self):
        return await self.page.content()

    # 页面滑动
    async def page_slide(self, level: int, vertical: int):
        """
        :param level: 正值表示向右滚动，负值表示向左滚动
        :param vertical: 正值表示向下滚动，负值表示向上滚动
        :param page: 指定页面，默认为当前页面
        """
        await self.page.evaluate(f"window.scrollBy({level}, {vertical});")

    # 元素滑动
    async def element_slide(self, selector: str, level: int, vertical: int):
        """
        :param selector: 元素定位
        :param level: 正值表示向右滚动，负值表示向左滚动
        :param vertical: 正值表示向下滚动，负值表示向上滚动
        :param page: 指定页面
        """
        page = self.page
        # 获取元素的定位器
        element = page.locator(selector).first

        # 获取元素的位置和大小
        bounding_box = await element.bounding_box()

        if bounding_box:
            # 计算起始位置
            start_x = bounding_box['x'] + bounding_box['width'] / 2
            start_y = bounding_box['y'] + bounding_box['height'] / 2

            # 计算目标位置
            end_x = start_x + level
            end_y = start_y + vertical

            # 移动鼠标到起始位置
            await page.mouse.move(start_x, start_y)
            # 按下鼠标
            await page.mouse.down()
            # 移动鼠标到目标位置
            await page.mouse.move(end_x, end_y)
            # 释放鼠标
            await page.mouse.up()
        else:
            print(f"未找到具有选择器“{selector}”的元素。")

    async def element_exist(self, selector):
        return await self.page.locator(selector).is_visible()

    async def switch_to_page(self, index):
        """
        切换到指定索引的页面
        :param index: 页面索引（0为第一个页面）
        :return: 切换后的页面对象
        """
        pages = self.context.pages  # 获取所有页面
        if 0 <= index < len(pages):
            self.page = pages[index]
            return self.page
        raise IndexError("页面索引超出范围")

    # 浏览器后退
    async def go_back(self):
        """浏览器后退到上一页"""
        await self.page.go_back()
        await self.page.wait_for_load_state("networkidle")  # 等待页面加载完成

    # 浏览器前进 
    async def go_forward(self):
        """浏览器前进到下一页"""
        await self.page.go_forward()
        await self.page.wait_for_load_state("networkidle")  # 等待页面加载完成


def safe_globals():
    return {
        "__builtins__": {
            "abs": abs,
            "all": all,
            "any": any,
            "bin": bin,
            "bool": bool,
            "chr": chr,
            "dict": dict,
            "divmod": divmod,
            "enumerate": enumerate,
            "filter": filter,
            "float": float,
            "format": format,
            "hash": hash,
            "hex": hex,
            "id": id,
            "int": int,
            "isinstance": isinstance,
            "issubclass": issubclass,
            "len": len,
            "list": list,
            "map": map,
            "max": max,
            "min": min,
            "oct": oct,
            "ord": ord,
            "pow": pow,
            "print": print,
            "range": range,
            "repr": repr,
            "reversed": reversed,
            "round": round,
            "slice": slice,
            "sorted": sorted,
            "str": str,
            "sum": sum,
            "tuple": tuple,
            "type": type,
            "zip": zip,
        },
        "md5": lambda s: hashlib.md5(s.encode('utf-8')).hexdigest(),
        "time": time,
    }


def case_leven():
    return {
        'blocker': Severity.BLOCKER,
        'critical': Severity.CRITICAL,
        'normal': Severity.NORMAL,
        'minor': Severity.MINOR,
        'trivial': Severity.TRIVIAL
    }


# asyncio.run(main())
async def send_message_to_redis(redis, room_id, message, image_data=None):
    try:
        # 如果有图片数据，将其添加到消息中
        if image_data:
            message['image'] = image_data
        await redis.publish('websocket_channel', json.dumps({"room_id": str(room_id), "case_status": message}))
        logger.info('消息发送成功')
    except Exception as e:
        logger.error(f"发送消息错误: {e}")

async def main():
    # 连接到Redis服务器
    redis = await aioredis.from_url('redis://localhost', encoding="utf-8", decode_responses=True)
    # 读取图片并转换为Base64编码
    with open('path/to/image.jpg', 'rb') as image_file:
        image_data = base64.b64encode(image_file.read()).decode('utf-8')
    # 发送消息到Redis，包括图片数据
    await send_message_to_redis(redis, 123, {"status": "running", "message": "添加到后台任务"}, image_data=image_data)
    # 关闭Redis连接
    await redis.close()

# asyncio.run(main())
